from time import time
from maix import display, camera , image


class LPR:
    loc_model_path = '/root/models/loc.mud'
    reg_model_path = '/root/models/reg.mud'
    chars =[ "皖", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑",
                    "苏", "浙", "京", "闽", "赣", "鲁", "豫", "鄂", "湘" , "粤",
                    "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青" , "宁",
                    "新", "警", "学", "A"  , "B" ,  "C" ,  "D" ,  "E" ,  "F"  ,  "G",
                    "H" ,   "J" ,  "K" ,  "L" , "M" , "N" ,  "P" ,  "Q" ,  "R" , "S",
                    "T" ,  "U" ,  "V" , "W", "X"  , "Y" , "Z" , "0" , "1", "2", "3",
                    "4", "5", "6", "7", "8", "9", "-"]

    variances = [0.1, 0.2]
    steps = [8, 16, 32]
    min_sizes = [12, 24, 48, 96, 192, 320]

    def __init__(self) -> None:
        from maix import nn
        # 使用try-except来安全地加载模型
        self.loc_model = None
        self.reg_model = None
        
        try:
            # 尝试使用通用方式加载模型
            # 先检查模型文件是否存在
            import os
            if not os.path.exists(self.loc_model_path):
                print(f"车牌定位模型文件不存在: {self.loc_model_path}")
                return
                
            if not os.path.exists(self.reg_model_path):
                print(f"字符识别模型文件不存在: {self.reg_model_path}")
                return
            
            # 尝试加载模型（使用与yolo_camera.py类似的方式）
            # 由于这是特定的车牌识别模型，我们不知道具体的类名，所以使用通用方式
            self.loc_model = nn.load(self.loc_model_path)
            self.reg_model = nn.load(self.reg_model_path)
            
            print("车牌识别模型加载成功")
        except AttributeError as e:
            print(f"模型加载函数不存在: {e}")
            print("尝试使用替代方法加载模型...")
            # 如果nn.load不存在，尝试其他方式
            try:
                # 检查是否有其他加载方法
                if hasattr(nn, 'model'):
                    self.loc_model = nn.model.Model()
                    self.loc_model.load(self.loc_model_path)
                    self.reg_model = nn.model.Model()
                    self.reg_model.load(self.reg_model_path)
                else:
                    print("未找到可用的模型加载方法")
                    return
            except Exception as e2:
                print(f"替代模型加载方法也失败了: {e2}")
                return
        except Exception as e:
            print(f"模型加载失败: {e}")
            return

        from maix.nn import decoder
        try:
            self.loc_decoder = decoder.license_plate_location([224,224] , self.steps , self.min_sizes, self.variances)
            self.reg_decoder  = decoder.CTC((1,68,18))
        except Exception as e:
            print(f"解码器初始化失败: {e}")
            return
        
        # 添加标志位，确保对象正确初始化
        self.initialized = True

    def __del__(self):
        # 在MaixPy中，通常不需要手动删除模型对象
        pass

    def cal_fps(self ,start , end):
        one_second = 1
        one_flash = end - start
        fps = one_second / one_flash
        return  fps

    def  draw_fps(self,img , fps):
        img.draw_string(0, 0 ,'FPS :'+str(fps), scale=1,color=(255, 0, 255), thickness=1)

    def draw_string(self , img , x , y , string , color):
        img.draw_string( x , y , string ,color = color)

    def draw_paste(self , src ,dst):
        src.paste(dst , 0 , 0)

    def draw_rectangle(self,img, box):
        img.draw_rectangle(box[0], box[1], box[2], box[3],color=(230 ,230, 250), thickness=2)

    def draw_point(self,img,landmark):
        for i in range(4):
            x = landmark[2 * i ]
            y = landmark[2 * i + 1]
            img.draw_rectangle(x-2,y-2, x+2,y+2,color= (193 ,255 ,193), thickness =-1)

    def process(self,input):
        # 确保模型已正确加载
        if not hasattr(self, 'initialized') or not self.initialized or self.loc_model is None or self.reg_model is None:
            # 如果模型未加载，则直接返回，不进行处理
            return
            
        try:
            loc_out = self.loc_model.forward(input, quantize=1, layout = "chw") # retinaface decoder only support chw layout
            boxes , landmarks = self.loc_decoder.run(loc_out, nms = 0.2 ,score_thresh = 0.7 , outputs_shape =[[1,4,2058],[1,2,2058],[1,8,2058]])

            for i,box in enumerate(boxes):

                landmark = landmarks[i][:6]
                reg_in  = input.crop_affine(landmark , 94 , 24)
                reg_out = self.reg_model.forward(reg_in ,  quantize=1, layout = "chw")

                LP_number = self.reg_decoder.run(reg_out)
                string_LP = ''
                for id in LP_number:
                    string_LP += self.chars[id]

                self.draw_string(input , box[0], box[1] , string_LP  ,color=(225,0,0))
                self.draw_paste(input , reg_in)
                self.draw_rectangle(input,box)
                self.draw_point(input , landmarks[i])
        except Exception as e:
            print(f"处理过程中出错: {e}")

def main():
    from maix import display, camera , image
    # image.load_freetype("/home/res/Sans.ttf")  # 注释掉这行，因为该函数不存在
    try:
        app  = LPR()
        while True:
            img = camera.capture().resize(size=(224,224))
            app.process(img)
            display.show(img)
            # break
    except Exception as e:
        print(f"程序运行出错: {e}")

# print('maix.version', maix.version)
if __name__ == "__main__":
    main()